import torch
import torch.nn.functional as F
from torch import nn
import torchvision.models


#'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M']

def make_layers(in_channels, out_channels, conv_ks=3, conv_s=1, conv_pad=1, pool_ks=3, pool_s=2, batch_norm=False):
    layers = []
    conv2d = nn.Conv2d(in_channels, out_channels, kernel_size=conv_ks, stride=conv_s, padding=conv_pad)
    if batch_norm:
        layers += [conv2d, nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True)]
    else:
        layers += [conv2d, nn.ReLU(inplace=True)]

    layers += [nn.MaxPool2d(kernel_size=pool_ks, stride=pool_s)]
    return nn.Sequential(*layers)


class LandMarkNet(nn.Module):
    """
    https://ieeexplore.ieee.org/document/8819955
    An Automated and Accurate Spine Curve Analysis System.
    """
    def __init__(self, n_inputs=3, n_outputs=68,
                 layer_channles={'layer1': 32, 'layer2': 64, 'layer3': 128, 'layer4': 256, 'layer5': 512},
                 image_shape= {'h': 512, 'w': 256}, batch_norm=False, init_weights=True):
        super(LandMarkNet, self).__init__()
        self.kp_num = n_outputs
        self.features = []

        self.features.append(make_layers(n_inputs, layer_channles['layer1'], conv_ks=7, conv_s=1, conv_pad=3, pool_ks=2, pool_s=2,
                                  batch_norm=batch_norm))
        self.features.append(make_layers(layer_channles['layer1'], layer_channles['layer2'], conv_ks=3, conv_s=1, conv_pad=1, pool_ks=2, pool_s=2,
                                  batch_norm=batch_norm))
        self.features.append(make_layers(layer_channles['layer2'], layer_channles['layer3'], conv_ks=3, conv_s=1, conv_pad=1, pool_ks=2, pool_s=2,
                                  batch_norm=batch_norm))
        self.features.append(make_layers(layer_channles['layer3'], layer_channles['layer4'], conv_ks=5, conv_s=1, conv_pad=2, pool_ks=2, pool_s=2,
                                  batch_norm=batch_norm))

        self.features = nn.Sequential(*self.features)
        self.classifier = nn.Sequential(
            nn.Linear(image_shape['h']//16 * image_shape['w']//16 *layer_channles['layer4'], layer_channles['layer5']),
            nn.ReLU(True),
            nn.Linear(layer_channles['layer5'], layer_channles['layer5']),
            nn.ReLU(True),
            nn.Linear(layer_channles['layer5'], n_outputs*2),
        )
        if init_weights:
            self._initialize_weights()

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        x = x.view(x.size(0), self.kp_num, 2)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)


def main():
    image_shape = (256, 128)
    a = torch.rand(2,3, image_shape[0], image_shape[1]).cuda()
    model = LandMarkNet(n_inputs=3, n_outputs=68,
                        layer_channles={'layer1': 32, 'layer2': 64, 'layer3': 128, 'layer4': 256, 'layer5': 512},
                        image_shape= {'h': image_shape[0], 'w': image_shape[1]}, batch_norm=False, init_weights=True)
    model = model.cuda()
    b = model.forward(a)

    print(b.shape)


if __name__ == '__main__':
    main()